package com.zjh.rocketmq.adapter;

import com.zjh.rocketmq.annotation.RocketmqTransactionListener;
import org.apache.rocketmq.spring.annotation.RocketMQTransactionListener;
import org.apache.rocketmq.spring.core.RocketMQLocalTransactionListener;
import org.apache.rocketmq.spring.core.RocketMQLocalTransactionState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.messaging.Message;
import org.springframework.stereotype.Component;

import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.util.stream.Stream;

@Component
@RocketMQTransactionListener(corePoolSize = 2, maximumPoolSize = 10)
public class RocketmqTransactionListenerAdapter implements RocketMQLocalTransactionListener, InitializingBean {
    private final Logger logger = LoggerFactory.getLogger(RocketmqTransactionListenerAdapter.class);
    private final ApplicationContext applicationContext;
    private final ConcurrentHashMap<String, RocketMQLocalTransactionListener> listeners;

    public RocketmqTransactionListenerAdapter(ApplicationContext applicationContext) {
        this.applicationContext = applicationContext;
        this.listeners = new ConcurrentHashMap<>();
    }
    @Override
    public RocketMQLocalTransactionState executeLocalTransaction(Message message, Object o) {
        String destination = getDestination(message);
        if (!listeners.containsKey(destination)) {
            logger.info("----------------------------------------------------");
            logger.warn("{} transaction message is not supported", destination);
            return RocketMQLocalTransactionState.ROLLBACK;
        } else {
            RocketMQLocalTransactionListener listener = listeners.get(destination);
            return listener.executeLocalTransaction(message, o);
        }
    }

    @Override
    public RocketMQLocalTransactionState checkLocalTransaction(Message message) {
        String destination = getDestination(message);
        if (!listeners.containsKey(destination)) {
            logger.info("----------------------------------------------------");
            logger.warn("{} transaction message is not supported", destination);
            return RocketMQLocalTransactionState.ROLLBACK;
        } else {
            RocketMQLocalTransactionListener listener = listeners.get(destination);
            return listener.checkLocalTransaction(message);
        }
    }

    private String getDestination(Message msg) {
        String topic = msg.getHeaders().get("rocketmq_TOPIC", String.class);
        String tags = msg.getHeaders().get("rocketmq_TAGS", String.class);
        return String.format("%s:%s", topic, tags);
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        applicationContext.getBeansWithAnnotation(RocketmqTransactionListener.class)
                .values()
                .stream()
                .filter(m -> RocketMQLocalTransactionListener.class.isAssignableFrom(m.getClass()))
                .collect(Collectors.toList())
                .forEach(listener -> {
                    RocketmqTransactionListener annotation =
                            listener.getClass().getAnnotation(RocketmqTransactionListener.class);

                    Stream.of(annotation.selectorExpression()
                                    .split("\\|\\|"))
                            .forEach(tag -> {
                                String destination = String.format("%s:%s", annotation.topic(), tag);
                                listeners.put(destination, (RocketMQLocalTransactionListener) listener);
                            });
                });
    }
}
